package com.example.gpu;

import jcuda.*;
import jcuda.driver.*;

import static jcuda.driver.JCudaDriver.*;

/**
 * CUDA工具类，提供一些常用的CUDA操作封装
 */
public class CudaUtils {
    
    /**
     * 初始化CUDA驱动并返回一个设备上下文
     * @param deviceId 设备ID，通常为0表示第一个GPU
     * @return CUDA上下文
     */
    public static CUcontext initCUDA(int deviceId) {
        try {
            // 初始化CUDA驱动
            cuInit(0);
            
            // 获取设备数量
            int[] deviceCount = new int[1];
            cuDeviceGetCount(deviceCount);
            
            if (deviceCount[0] == 0) {
                System.out.println("未检测到NVIDIA GPU设备");
                return null;
            }
            
            // 检查设备ID是否有效
            if (deviceId < 0 || deviceId >= deviceCount[0]) {
                System.out.println("无效的设备ID，使用默认设备 0");
                deviceId = 0;
            }
            
            // 获取指定的GPU设备
            CUdevice device = new CUdevice();
            cuDeviceGet(device, deviceId);
            
            // 获取设备名称
            byte[] deviceName = new byte[100];
            cuDeviceGetName(deviceName, deviceName.length, device);
            System.out.println("使用GPU设备: " + new String(deviceName).trim());
            
            // 获取设备计算能力
            int[] computeCapability = new int[2];
            cuDeviceGetAttribute(computeCapability, CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR, device);
            cuDeviceGetAttribute(computeCapability, CUdevice_attribute.CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR, device);
            System.out.println("计算能力: " + computeCapability[0] + "." + computeCapability[1]);
            
            // 创建CUDA上下文
            CUcontext context = new CUcontext();
            cuCtxCreate(context, 0, device);
            
            return context;
        } catch (Exception e) {
            System.err.println("CUDA初始化失败: " + e.getMessage());
            e.printStackTrace();
            return null;
        }
    }
    
    /**
     * 释放CUDA资源
     * @param context CUDA上下文
     * @param module CUDA模块
     * @param devicePointers 设备内存指针数组
     */
    public static void releaseCUDA(CUcontext context, CUmodule module, CUdeviceptr... devicePointers) {
        // 释放设备内存
        if (devicePointers != null) {
            for (CUdeviceptr ptr : devicePointers) {
                if (ptr != null) {
                    cuMemFree(ptr);
                }
            }
        }
        
        // 卸载模块
        if (module != null) {
            cuModuleUnload(module);
        }
        
        // 销毁上下文
        if (context != null) {
            cuCtxDestroy(context);
        }
    }
    
    /**
     * 从PTX文件加载CUDA模块
     * @param ptxFilePath PTX文件路径
     * @param kernelName 内核函数名称
     * @return 包含模块和函数的数组 [module, function]
     */
    public static Object[] loadKernelFromPTX(String ptxFilePath, String kernelName) {
        try {
            CUmodule module = new CUmodule();
            cuModuleLoad(module, ptxFilePath);
            
            CUfunction function = new CUfunction();
            cuModuleGetFunction(function, module, kernelName);
            
            return new Object[] {module, function};
        } catch (Exception e) {
            System.err.println("加载CUDA内核失败: " + e.getMessage());
            e.printStackTrace();
            return null;
        }
    }
    
    /**
     * 将float数组从主机内存复制到设备内存
     * @param hostArray 主机数组
     * @return 设备内存指针
     */
    public static CUdeviceptr copyToDevice(float[] hostArray) {
        if (hostArray == null) {
            return null;
        }
        
        CUdeviceptr devicePtr = new CUdeviceptr();
        cuMemAlloc(devicePtr, hostArray.length * Sizeof.FLOAT);
        cuMemcpyHtoD(devicePtr, Pointer.to(hostArray), hostArray.length * Sizeof.FLOAT);
        
        return devicePtr;
    }
    
    /**
     * 将float数组从设备内存复制到主机内存
     * @param devicePtr 设备内存指针
     * @param hostArray 主机数组
     */
    public static void copyFromDevice(CUdeviceptr devicePtr, float[] hostArray) {
        if (devicePtr == null || hostArray == null) {
            return;
        }
        
        cuMemcpyDtoH(Pointer.to(hostArray), devicePtr, hostArray.length * Sizeof.FLOAT);
    }
    
    /**
     * 计算合适的网格大小
     * @param n 元素数量
     * @param blockSize 块大小
     * @return 网格大小
     */
    public static int calculateGridSize(int n, int blockSize) {
        return (n + blockSize - 1) / blockSize;
    }
    
    /**
     * 检查CUDA错误
     * @param errorCode CUDA错误码
     * @param message 错误消息
     * @return 是否有错误
     */
    public static boolean checkCudaError(int errorCode, String message) {
        if (errorCode != CUresult.CUDA_SUCCESS) {
            int errorString = cuGetErrorString(errorCode, null);
            System.err.println(message + ": " + errorString);
            return true;
        }
        return false;
    }
}